/*
Copyright (c) 2025 Huawei Technologies Co., Ltd.
This file is a part of the CANN Open Software.
Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
Please refer to the License for details. You may not use this file except in compliance with the License.
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
See LICENSE in the root of the software repository for the full text of the License. 
*/

#pragma once
#include "acl/acl.h"
#include "../include/custom_type.h"
#include "./data_utils.h"
#include <cstdio>
#include <string>
#include <algorithm>

/**
 * @brief: main函数中调用host的工具函数
 * @param [in] kernel_host: 本次调用的host函数
 * @param [in] layoutA: A矩阵排布格式
 * @param [in] layoutB: B矩阵排布格式
 * @param [in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param [in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param [in] zeroPaddingK: A、B矩阵零填充后的K维度
 * @param [in] batchCount: 批量矩阵乘的batch数
 * @param [in] pathA: 本次GEMM计算的A矩阵数据（如果有）读取路径
 * @param [in] pathB: 本次GEMM计算的B矩阵数据（如果有）读取路径
 * @param [in] pathC: 本次GEMM计算的C矩阵数据（如果有）读取路径
 * @param [in] pathAlpha: 本次GEMM计算的alpha数据（如果有）读取路径
 * @param [in] pathBeta:  本次GEMM计算的beta 数据（如果有）读取路径
 * @param [in] pathMaskA: 本次GEMM计算的maskM数据（如果有）读取路径
 * @param [in] pathExpectResult: 本次GEMM计算的golden数据（如果有）读取路径
 * @param [in] preOutput: 上一任务的输出（如果有），作为本次任务的输入，最后会free掉
 * @param [out] curOutput: 本次任务的输出，作为下一任务的输入（如果有），最后会返回
 * @param [in] mode: 0-strideBatchedGemm的满维度矩阵计算，123-三个LLMsGEMM任务有效矩阵维度计算
 * @param [in] testMode: 是否需要读取和验证真值数据。0-功能测试，1-性能测试
 */
void run_main(
    void (* kernel_host)(
        layoutType,         // layoutA
        layoutType,         // layoutB
        uint32_t,           // zeroPaddingM
        uint32_t,           // zeroPaddingN
        uint32_t,           // zeroPaddingK
        uint32_t,           // batchCount
        half *,             // d_maskA
        half,               // alpha
        half *,             // d_A
        half *,             // d_B
        half,               // beta
        half *,             // d_C
        aclrtStream &,      // stream
        uint8_t             // mode
    ), 
    aclrtStream &stream, 
    layoutType layoutA, 
    layoutType layoutB, 
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t zeroPaddingK, 
    uint32_t batchCount, 
    std::string &pathA, 
    std::string &pathB, 
    std::string &pathC, 
    std::string &pathAlpha, 
    std::string &pathBeta, 
    std::string &pathMaskA,  
    std::string &pathExpectResult, 
    half* &preOutput,          
    half* &curOutput, 
    uint8_t mode, 
    uint8_t testMode
){

    INFO_LOG("Main method starts execution. ");

    if(testMode == 0){ 

        uint8_t printMatrix = 0; 
        uint8_t printAlphaBeta = 0; 
        uint8_t printResult = 0; 

        uint64_t inputABytes = zeroPaddingM * zeroPaddingK * sizeof( half ) * batchCount;
        uint64_t inputBBytes = zeroPaddingK * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputCBytes = zeroPaddingM * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputAlphaBetaBytes = sizeof( half );
        uint64_t inputMaskABytes= sizeof(half) * zeroPaddingM * batchCount; 

        half *h_A = nullptr;
        half *h_B = nullptr;
        half *h_C = nullptr;
        half *h_alpha = nullptr;
        half *h_beta = nullptr;
        half *h_maskA = nullptr; 

        if(preOutput == nullptr)
            ACL_CHECK(aclrtMallocHost((void **)&h_A, inputABytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_B, inputBBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_C, inputCBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_alpha, inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_beta , inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_maskA, inputMaskABytes));

        if(preOutput == nullptr)
            ReadFile( pathA, h_A, inputABytes );
        ReadFile( pathB, h_B, inputBBytes );
        ReadFile( pathC, h_C, inputCBytes );
        ReadFile( pathAlpha, h_alpha, inputAlphaBetaBytes );
        ReadFile( pathBeta , h_beta , inputAlphaBetaBytes );
        ReadFile( pathMaskA, h_maskA, inputMaskABytes );

        if(printMatrix){

            printf("A:\n");
            for(uint32_t i = 0; i < batchCount * zeroPaddingM; i++ ){
                for(uint32_t j = 0; j < zeroPaddingK; j++){
                    printf("[%d, %d]%f ", i, j, h_A[i*zeroPaddingK+j]);
                }
                printf("\n");
            }
            printf("B:\n");
            for(uint32_t i = 0; i < batchCount * zeroPaddingK; i++ ){
                for(uint32_t j = 0; j < zeroPaddingN; j++){
                    printf("[%d, %d]%f ", i, j, h_B[i*zeroPaddingN+j]);
                }
                printf("\n");
            }
            printf("C:\n");
            for(uint32_t i = 0; i < batchCount * zeroPaddingM; i++ ){
                for(uint32_t j = 0; j < zeroPaddingN; j++){
                    printf("[%d, %d]%f ", i, j, h_C[i*zeroPaddingN+j]);
                }
                printf("\n");
            }
        }

        if(printAlphaBeta){
            printf("alpha = %f, beta = %f\n", *h_alpha, *h_beta);
        }

        half *d_A = nullptr;
        half *d_B = nullptr;
        half *d_C = nullptr;
        half *d_maskA = nullptr; 
        half alpha = 0.0;
        half beta  = 0.0;

        if(preOutput == nullptr)
            ACL_CHECK(aclrtMalloc((void **)&d_A, inputABytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_B, inputBBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_C, inputCBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_maskA, inputMaskABytes, ACL_MEM_MALLOC_HUGE_FIRST));

        alpha = *h_alpha;
        beta  = *h_beta;
        if(preOutput == nullptr)
            ACL_CHECK(aclrtMemcpy(d_A, inputABytes, h_A, inputABytes, ACL_MEMCPY_HOST_TO_DEVICE));
        else
            d_A = preOutput;
        ACL_CHECK(aclrtMemcpy(d_B, inputBBytes, h_B, inputBBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_C, inputCBytes, h_C, inputCBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_maskA, inputMaskABytes, h_maskA, inputMaskABytes, ACL_MEMCPY_HOST_TO_DEVICE));

        (*kernel_host)(
            layoutA, 
            layoutB, 
            zeroPaddingM, 
            zeroPaddingN, 
            zeroPaddingK, 
            batchCount, 
            d_maskA, 
            alpha, 
            d_A, 
            d_B, 
            beta, 
            d_C, 
            stream, 
            mode
        );

        ACL_CHECK(aclrtMemcpy(h_C, inputCBytes, d_C, inputCBytes, ACL_MEMCPY_DEVICE_TO_HOST));

        half *h_expect_result = nullptr;
        ACL_CHECK(aclrtMallocHost((void **)&h_expect_result, inputCBytes));
        ReadFile( pathExpectResult, h_expect_result, inputCBytes );

        if(printResult){

            printf("expect_result:\n");
            for(uint32_t i = 0; i < batchCount * zeroPaddingM; i++ ){
                for(uint32_t j = 0; j < zeroPaddingN; j++){
                    printf("[%d, %d]%f ", i, j, h_expect_result[i*zeroPaddingN+j]);
                }
                printf("\n");
            }

            printf("output_C:\n");
            for(uint32_t i = 0; i < batchCount * zeroPaddingM; i++ ){
                for(uint32_t j = 0; j < zeroPaddingN; j++){
                    printf("[%d, %d]%f ", i, j, h_C[i*zeroPaddingN+j]);
                }
                printf("\n");
            }

        }

        uint32_t errorCount = calOutputError<half>( h_C, h_expect_result, zeroPaddingM * zeroPaddingN * batchCount, zeroPaddingK, zeroPaddingM, zeroPaddingN );

        if(preOutput == nullptr)
            ACL_CHECK(aclrtFreeHost(h_A));
        ACL_CHECK(aclrtFreeHost(h_B));
        ACL_CHECK(aclrtFreeHost(h_C));
        ACL_CHECK(aclrtFreeHost(h_alpha));
        ACL_CHECK(aclrtFreeHost(h_beta));
        ACL_CHECK(aclrtFreeHost(h_maskA));
        ACL_CHECK(aclrtFreeHost(h_expect_result));

        ACL_CHECK(aclrtFree(d_A));
        ACL_CHECK(aclrtFree(d_B));
        curOutput = d_C;
        ACL_CHECK(aclrtFree(d_maskA));

        INFO_LOG("Main method has been executed successfully. ErrorCount = %d. ", errorCount);

    }else if(testMode == 1){

        uint64_t inputABytes = zeroPaddingM * zeroPaddingK * sizeof( half ) * batchCount;
        uint64_t inputBBytes = zeroPaddingK * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputCBytes = zeroPaddingM * zeroPaddingN * sizeof( half ) * batchCount;
        uint64_t inputAlphaBetaBytes = sizeof( half );
        uint64_t inputMaskABytes = sizeof( half ) * batchCount * zeroPaddingM; 

        half *h_A = nullptr;
        half *h_B = nullptr;
        half *h_C = nullptr;
        half *h_alpha = nullptr;
        half *h_beta = nullptr;
        half *h_maskA = nullptr; 

        if(preOutput == nullptr)
            ACL_CHECK(aclrtMallocHost((void **)&h_A, inputABytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_B, inputBBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_C, inputCBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_alpha, inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_beta , inputAlphaBetaBytes));
        ACL_CHECK(aclrtMallocHost((void **)&h_maskA, inputMaskABytes));

        half *d_A = nullptr;
        half *d_B = nullptr;
        half *d_C = nullptr;
        half *d_maskA = nullptr; 
        half alpha = 0.0; 
        half beta  = 0.0; 

        if(preOutput == nullptr)
            ACL_CHECK(aclrtMalloc((void **)&d_A, inputABytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_B, inputBBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_C, inputCBytes, ACL_MEM_MALLOC_HUGE_FIRST));
        ACL_CHECK(aclrtMalloc((void **)&d_maskA, inputMaskABytes, ACL_MEM_MALLOC_HUGE_FIRST));

        alpha = 0.5f;
        beta  = 0.5f;
        
        printf("\n");
        printf("alpha = %f, beta = %f\n", alpha, beta);
        printf("\n");

        if(preOutput == nullptr)
            ACL_CHECK(aclrtMemcpy(d_A, inputABytes, h_A, inputABytes, ACL_MEMCPY_HOST_TO_DEVICE));
        else
            d_A = preOutput;
        ACL_CHECK(aclrtMemcpy(d_B, inputBBytes, h_B, inputBBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_C, inputCBytes, h_C, inputCBytes, ACL_MEMCPY_HOST_TO_DEVICE));
        ACL_CHECK(aclrtMemcpy(d_maskA, inputMaskABytes, h_maskA, inputMaskABytes, ACL_MEMCPY_HOST_TO_DEVICE));

        (*kernel_host)(
            layoutA, 
            layoutB, 
            zeroPaddingM, 
            zeroPaddingN, 
            zeroPaddingK, 
            batchCount, 
            d_maskA, 
            alpha, 
            d_A, 
            d_B, 
            beta, 
            d_C, 
            stream, 
            0
        );

        if(preOutput == nullptr)
            ACL_CHECK(aclrtFreeHost(h_A));
        ACL_CHECK(aclrtFreeHost(h_B));
        ACL_CHECK(aclrtFreeHost(h_C));
        ACL_CHECK(aclrtFreeHost(h_alpha));
        ACL_CHECK(aclrtFreeHost(h_beta));
        ACL_CHECK(aclrtFreeHost(h_maskA));

        ACL_CHECK(aclrtFree(d_A));
        ACL_CHECK(aclrtFree(d_B));
        curOutput = d_C;
        ACL_CHECK(aclrtFree(d_maskA));

        INFO_LOG("Main method has been executed successfully. ");

    }

    

}